import os

os.environ["CUDA_VISIBLE_DEVICES"] = '2,3'
from langchain.agents import create_sql_agent
from langchain.agents.agent_toolkits import SQLDatabaseToolkit
from langchain.agents.agent_types import AgentType
from langchain.sql_database import SQLDatabase

from langchain.llms import VLLM
from langchain import OpenAI, SQLDatabase
from langchain_experimental.sql import SQLDatabaseChain
from langchain.chains import create_sql_query_chain
import jsonlines
import traceback
from modelscope import AutoModelForCausalLM, AutoTokenizer, snapshot_download
from modelscope import GenerationConfig

model_dir = '/datasets/fengjiahao/nlp/TongyiFinance/Tongyi-Finance-14B'
# model_dir = '/datasets/fengjiahao/nlp/TongyiFinance/Tongyi-Finance-14B-Chat-Int4/'
question_json_path = r'/datasets/fengjiahao/nlp/bs_challenge_financial_14b_dataset/question.json'
answer_path = r'/datasets/fengjiahao/nlp/bs_challenge_financial_14b_dataset/submit_result.jsonl'
content = []
with jsonlines.open(question_json_path, "r") as json_file:
    for obj in json_file.iter(type=dict, skip_invalid=True):
        content.append(obj)
llm = VLLM(
    model=model_dir,
    trust_remote_code=True,  # mandatory for hf models
    temperature=0.2,
    top_p=0.7,
    top_k=10,
    tensor_parallel_size=2, verbose=True
)
# llm = AutoModelForCausalLM.from_pretrained(model_dir, device_map="auto", trust_remote_code=True).eval()
# llm.generation_config = GenerationConfig.from_pretrained(model_dir, trust_remote_code=True)
table_dict = {
    # 'A股': ['A股票日行情表', 'A股股票行业划分表', 'A股股票行情'],
    'A股': [ 'A股股票行情'],
    '基金': ['基金份额持有人结构', '基金债券持仓明细', '基金可转债持仓明细', '基金基本信息', '基金日行情表', '基金股票持仓明细', '基金规模变动表'],
    '港股': ['港股票日行情表'],
}


def ask_llm(question):
    theme = 'A股'
    if '基金' in question:
        theme = '基金'
    elif '港股' in question or '香港' in question:
        theme = '港股'

    target_table_list = table_dict[theme]
    db = SQLDatabase.from_uri(
        "sqlite:////datasets/fengjiahao/nlp/bs_challenge_financial_14b_dataset/dataset/博金杯比赛数据.db",
        include_tables=target_table_list,
        view_support=True,
        sample_rows_in_table_info=2)
    toolkit = SQLDatabaseToolkit(db=db, llm=llm)
    agent_executor = create_sql_agent(
        llm=llm,
        toolkit=toolkit,
        verbose=True,
        agent_type=AgentType.ZERO_SHOT_REACT_DESCRIPTION,
        max_iterations=5,
    )
    print(db.table_info)
    try:
        response = agent_executor.run(question)
    except Exception as e:
        traceback.print_exc()
        print(e)
        response = ''
    print("!!!Q:", question)
    print("!!!A:", response)
    return response


for cont in content:
    question = cont['question']
    response = ask_llm(question)
    # print(question)

    cont['answer'] = response
    # 市盈率是最常用来评估股价水平是否合理的指标之一，是指股票价格与每股盈利的比率。...
with jsonlines.open(answer_path, "w") as json_file:
    json_file.write_all(content)
